This pull request includes multiple changes to improve code readability, simplify function definitions, and update method calls across various files in the brainstate package. The most important changes include refactoring functions, cleaning up imports, and updating progress bar handling.

close #44

Code readability improvements:

• brainstate/compile/_loop_collect_return.py: Refactored the scan function to improve readability by formatting the parameters of jax.lax.scan across multiple lines.

Function and method updates:

• brainstate/compile/_make_jaxpr.py: Refactored _init_state_trace to _init_state_trace_stack, including a new method _new_jax_trace to handle different JAX versions, and updated calls to use the new method. [1] [2] [3]

Import cleanup:

• brainstate/compile/_make_jaxpr.py: Removed an unused import of new_jax_trace.
• brainstate/util/_tracers.py: Removed the new_jax_trace function and its related import, as it is now defined elsewhere. [1] [2]

Progress bar handling:

• brainstate/compile/_progress_bar.py: Simplified the progress bar update logic by introducing a new _tqdm method and updating the __call__ method to use jax.debug.callback.

Example updates:

• examples/100_hh_neuron_model.py: Improved readability by formatting the parameters of bst.compile.for_loop across multiple lines.
• examples/106_COBA_HH_2007.py: Removed an unused import of dendritex and updated the call to bst.nn.init_all_states to exclude the dx.IonChannel parameter. [1] [2]